/* Copyright 2013 Tobias Marschall
 *
 * This file is part of CLEVER.
 *
 * CLEVER is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * CLEVER is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with CLEVER.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <cassert>
#include <sstream>
#include <cmath>

#include "GenotypeDistribution.h"

using namespace std;

GenotypeDistribution::GenotypeDistribution(double no_prob, double hetero_prob, double homo_prob) : distribution() {
	distribution.push_back(no_prob);
	distribution.push_back(hetero_prob);
	distribution.push_back(homo_prob);
}

GenotypeDistribution::genotype_t GenotypeDistribution::likeliestGenotype() {
	int best_index = 0;
	double best = 0.0;
	for (size_t i=0; i<distribution.size(); ++i) {
		if (distribution[i] > best) {
			best = distribution[i];
			best_index = i;
		}
	}
	return (genotype_t)best_index;
}

std::string GenotypeDistribution::genotypeToString(genotype_t genotype) {
	switch (genotype) {
		case ABSENT: return "0/0";
		case HETEROZYGOUS: return "0/1";
		case HOMOZYGOUS: return "1/1";
		default: assert(false);
	}
}

string GenotypeDistribution::likeliestGenotypeString() {
	return genotypeToString(likeliestGenotype());
}

void GenotypeDistribution::normalize() {
	double p_sum = 0.0;
	for (size_t i=0; i<distribution.size(); ++i) {
		 p_sum += distribution[i];
	}
	if (p_sum <= 0.0) {
		distribution = vector<double>(3, 1.0/3.0);
	} else {
		for (size_t i=0; i<distribution.size(); ++i) {
			distribution[i] /= p_sum;
		}
	}
}

double GenotypeDistribution::errorProbability() {
	int best_index = 0;
	double best = 0.0;
	for (size_t i=0; i<distribution.size(); ++i) {
		if (distribution[i] > best) {
			best = distribution[i];
			best_index = i;
		}
	}
	double p = 0.0;
	for (size_t i=0; i<distribution.size(); ++i) {
		if (i == best_index) continue;
		p += distribution[i];
	}
	return p;
}

double GenotypeDistribution::errorProbability(genotype_t genotype) {
	double p = 0.0;
	for (size_t i=0; i<distribution.size(); ++i) {
		if (i == genotype) continue;
		p += distribution[i];
	}
	return p;
}

GenotypeDistribution operator*(const GenotypeDistribution& d1, const GenotypeDistribution& d2) {
	vector<double> d(d1.distribution);
	double sum = 0.0;
	for (int i=0; i<3; ++i) {
		d[i] *= d2.distribution[i];
		sum += d[i];
	}
	for (int i=0; i<3; ++i) d[i] /= sum;
	return GenotypeDistribution(d[0], d[1], d[2]);
}

string GenotypeDistribution::toPLString() {
	double max = 0.0;
	for (int i=0; i<3; ++i) {
		if (distribution[i] > max) max = distribution[i];
	}
	if (max == 0.0) return "0,0,0";
	ostringstream oss;
	for (int i=0; i<3; ++i) {
		if (i>0) oss << ',';
		oss << ((int)round(min(255.0,-log10(distribution[i]/max)*10.0)));
	}
	return oss.str();
}

ostream& operator<<(ostream& os, const GenotypeDistribution& d) {
	os << "(NO:" << d.notPresent() << ", HET:" << d.heterozygous() << ", HOMO:" << d.homozyguous() << ")";
	return os;
}

std::ostream& operator<<(std::ostream& os, const GenotypeDistribution::genotype_t& g) {
	switch (g) {
		case GenotypeDistribution::ABSENT: 
			os << "0/0";
			break;
		case GenotypeDistribution::HETEROZYGOUS:
			os << "0/1";
			break;
		case GenotypeDistribution::HOMOZYGOUS:
			os << "1/1";
			break;
		default:
			assert(false);
	}
}
